from itertools import izip,count
from astro.utilities import readatom
from astro import spec
import numpy as np
import pylab as pl

# Pixel Optical Depth algorithm

# Note a low nsigcont will include more dodgy points that will stuff
# up the N(HI) estimates, but tend to show up metal absorption
# better. Vice versa for a high nsigcont value. 

# for metal finding you probably only need a few orders, say 5, and a
# low nsigcont, say 0.5 or 1.

def calcpod(spec, z=-1., minwav=-1., atomfilename='/home/vpfit/atom.dat',
            n_ord=50, nsigsat=3, nsigcont=3):
    """ Input:

    Spectrum class object
    atom.dat filename with path
    QSO redshift (for finding wavs over which to measure Lya POD)
    Minimum useful wavelength for searching for higher order lines

    n_ord = Number of higher order transitions to use

    nsigsat = Reject points if they are less than this many sigma away from 
    zero flux (i.e. class them as saturated).

    nsigcont = Reject points if they are less than this many sigma away from 
    the continuum.


    Output:

    The optical depth using just the lya region
    Optical depth using all the higher order information
    """
    if z == -1.:
        raise Exception('Must supply the QSO redshift, z!')
    elif minwav == -1.:
        raise Exception('Must supply the cutoff wavelength, minwav!')

    # read in atom.dat data into dictionary
    atom = readatom(atomfilename)
    LYA_WAV = atom['HI'].wav[0]
    LYA_OSC = atom['HI'].osc[0]
    # [2:] excludes LL entry in atom.dat
    atom['HI'].sort(order='osc')
    HI_transitions = atom['HI'][2:]
    # reverse order (want highest -> lowest osc strength)
    HI_transitions = HI_transitions[::-1]
    HI_transitions = HI_transitions[:n_ord]
    print 'Using %s higher order transitions' % len(HI_transitions)
    #print HI_transitions
    minosc = HI_transitions.osc[-1]
    #print minosc

    # these will have to be adjusted...
    wavmax = min(1215.6701*(1.+z), spec.wa[-1])
    wavmin = max(1025.72*(1.+z), minwav)
    i1,i2 = spec.wa_index(wavmin),spec.wa_index(wavmax)

    podwav = []
    pod = []
    pod2 = []
    pts = []         # total available points (from higher order lines) 
    foo=[]

    # for each pixel in the HI lya region
    for i,flux in izip(count(i1), spec.fl[i1:i2+1]):
        podwav.append(spec.wa[i])

        ######################################################
        # for comparison, make a tau array only using lya info
        ######################################################

        flux,cont,sig = spec.fl[i],spec.co[i],spec.er[i]
        if cont != -1. and sig > 0.0:
            flux,sig = flux/cont, sig/cont
            if flux < nsigsat*sig:    # if line saturated:
                pod2.append(-np.log(nsigsat*sig))
            elif flux > 1:             # if negative tau value:
                pod2.append(0.0)
            else:
                pod2.append(-np.log(flux))
        else:
            pod2.append(-1.)
 
        ######################################################
        # Real tau measurement
        ######################################################

        tau = [];  tausat = [];  tau_thin = []

        # Look at tau for every possible HI order transition
        for wav,osc,gam in HI_transitions:
            sf = LYA_WAV * LYA_OSC / (wav * osc) # scale factor
            zp1 = spec.wa[i] / LYA_WAV
            obswav = wav * zp1
            # check this transition is still inside the wav range of spectrum!
            if obswav < minwav:
                break
            j = spec.wa_index(obswav)
            flux,cont,sig = spec.fl[j],spec.co[j],spec.er[j]
            if cont == -1. or sig <= 0.0:   # bad pixel
                continue
            # Switch to normalised flux, sigma values
            flux,sig = flux / cont, sig / cont
            if nsigsat*sig < flux < 1. - nsigcont*sig:   # then good pixel
                tau.append(-np.log(flux) * sf)
                continue
            if flux >= 1 - nsigcont*sig and nsigcont*sig < 1.:
                # value close to continuum: ignore unless it's Lya.
                if sf == 1.:     # if lya
                    tau.append(max(0.,-np.log(flux) * sf))
                tau_thin.append(max(0.,-np.log(flux) * sf))
                continue
            if flux <= nsigsat*sig and nsigsat*sig < 1.:
                # saturated: Ignore unless no good pixels
                tausat.append(-np.log(nsigsat * sig) * sf)
            # If you get here the pixel has no contraint on the optiacl
            # depth, skip
            continue

        if len(tau) > 0:
            foo.append(tau)     # keep record of tau values used for this pixel
            pod.append(min(tau))
            pts.append(len(tau)) 
        # if no good pixels, see if all remaining pixels were
        # consistent with being saturated pixels
        elif len(tausat) > 0 and len(tau_thin) == 0:
            # assume it's a high tau pixel (i.e. we assume that
            # it's not all randomly placed, unrelated absorption -
            # this assumption will be poorer with smaller numbers of
            # transitions used!)
            foo.append(tausat)  # keep record of tau values used
            pod.append(max(tausat)) # find the maximum tau
            pts.append(len(tausat))
        else:
            foo.append([])
            pod.append(-1.)
            pts.append(0)
 
    lya = np.rec.fromarrays([pod2,podwav],names='tau,wa')
    all = np.rec.fromarrays([pod,pts,podwav],names='tau,pts,wa')

    return lya,all,foo

def plotpod(lya,all):
    f = pl.figure(figsize=(14,4))
    #f.subplots_adjust(left=0.05,right=0.97)
    a = f.add_subplot(111)
    a.plot(all.wa,all.pts,'g',ls='steps')
    a.plot(all.wa,all.tau,'r')
    a.plot(lya.wa,lya.tau,'b')
    a.legend(['no. pixels','All','Lya'])
    a.set_ylabel('Optical depth')
    a.set_xlabel('Wavelength (Ang)')
    pl.show()

if 0:
    z = 3.819
    # Don't look for higher order lyman series past this wavelength
    wavcutoff = 4172.
    sp = spec.read('/home/nhmc/data/uves/J0124+0044/J0124p0044.fits')
    lya,all,foo = calcpod(sp,z=z,minwav=wavcutoff,
                          nsigcont=3.,nsigsat=3.,n_ord=30)
    plotpod(lya,all)
    
